"""
PySide6 TableView 多级表头示例
演示如何实现多行多列合并的表头
"""

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QTableView, QVBoxLayout, 
    QWidget, QHeaderView, QStyleOptionHeader, QStyle
)
from PySide6.QtGui import QStandardItemModel, QStandardItem, QPainter, QColor
from PySide6.QtCore import Qt, QRect, QSize
import sys


class MultiLevelHeaderView(QHeaderView):
    """多级表头视图"""
    
    def __init__(self, orientation, levels, parent=None):
        """
        初始化多级表头
        
        参数:
            orientation: Qt.Orientation.Horizontal 或 Qt.Orientation.Vertical
            levels: 表头层级定义，格式:
                [
                    [(text, col_start, col_span), ...],  # 第1行
                    [(text, col_start, col_span), ...],  # 第2行
                    ...
                ]
        """
        super().__init__(orientation, parent)
        self._levels = levels
        self._level_count = len(levels)
        self.setSectionsClickable(False)  # 禁用点击排序
        
    def sizeHint(self):
        """返回表头的建议大小（高度为层级数 * 单行高度）"""
        base_size = super().sizeHint()
        if self.orientation() == Qt.Orientation.Horizontal:
            # 水平表头：高度为层级数 * 30
            return QSize(base_size.width(), self._level_count * 30)
        else:
            return base_size
    
    def paintEvent(self, event):
        """重写paintEvent来绘制整个表头"""
        painter = QPainter(self.viewport())
        painter.save()
        
        # 绘制背景
        painter.fillRect(self.viewport().rect(), QColor(240, 240, 240))
        
        level_height = self.viewport().height() // self._level_count
        
        # 绘制每一层级
        for level_idx, level_items in enumerate(self._levels):
            y_offset = level_idx * level_height
            
            for text, col_start, col_span in level_items:
                # 计算合并单元格的位置和大小
                x_start = self.sectionViewportPosition(col_start)
                width = sum(self.sectionSize(col_start + i) for i in range(col_span))
                
                cell_rect = QRect(x_start, y_offset, width, level_height)
                
                # 绘制边框
                painter.setPen(QColor(180, 180, 180))
                painter.drawRect(cell_rect.adjusted(0, 0, -1, -1))
                
                # 绘制文本
                painter.setPen(QColor(0, 0, 0))
                painter.drawText(cell_rect, Qt.AlignmentFlag.AlignCenter, text)
        
        painter.restore()


class MultiHeaderDemo(QMainWindow):
    """多级表头演示主窗口"""
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.setup_model()
        self.setup_multi_header()
    
    def init_ui(self):
        """初始化UI"""
        self.setWindowTitle('TableView 多级表头示例')
        self.setGeometry(100, 100, 900, 600)
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # TableView
        self.table_view = QTableView()
        self.table_view.setAlternatingRowColors(True)
        layout.addWidget(self.table_view)
    
    def setup_model(self):
        """设置数据模型"""
        self.model = QStandardItemModel(5, 9)
        
        # 设置底层表头（最细粒度的列名）
        headers = ['姓名', '性别', '年龄', '语文', '数学', '英语', '物理', '化学', '生物']
        self.model.setHorizontalHeaderLabels(headers)
        
        # 填充示例数据
        data = [
            ['张三', '男', '18', '85', '90', '88', '92', '87', '89'],
            ['李四', '女', '17', '92', '88', '91', '85', '90', '88'],
            ['王五', '男', '18', '78', '85', '82', '88', '84', '86'],
            ['赵六', '女', '17', '88', '92', '90', '91', '89', '92'],
            ['孙七', '男', '18', '90', '87', '89', '86', '88', '85'],
        ]
        
        for row_idx, row_data in enumerate(data):
            for col_idx, value in enumerate(row_data):
                item = QStandardItem(value)
                item.setTextAlignment(Qt.AlignmentFlag.AlignCenter)
                self.model.setItem(row_idx, col_idx, item)
        
        self.table_view.setModel(self.model)
    
    def setup_multi_header(self):
        """设置多级表头"""
        # 定义三级表头结构
        # 格式：[(显示文本, 起始列, 跨列数), ...]
        levels = [
            # 第1层：最顶层分类
            [
                ('基本信息', 0, 3),      # 跨3列：姓名、性别、年龄
                ('成绩', 3, 6),          # 跨6列：所有科目
            ],
            # 第2层：中间分类
            [
                ('个人', 0, 1),          # 姓名
                ('基础', 1, 2),          # 性别、年龄
                ('文科', 3, 3),          # 语文、数学、英语
                ('理科', 6, 3),          # 物理、化学、生物
            ],
            # 第3层：最底层（具体列名）
            [
                ('姓名', 0, 1),
                ('性别', 1, 1),
                ('年龄', 2, 1),
                ('语文', 3, 1),
                ('数学', 4, 1),
                ('英语', 5, 1),
                ('物理', 6, 1),
                ('化学', 7, 1),
                ('生物', 8, 1),
            ],
        ]
        
        # 创建并设置多级水平表头
        multi_header = MultiLevelHeaderView(Qt.Orientation.Horizontal, levels, self.table_view)
        self.table_view.setHorizontalHeader(multi_header)
        
        # 调整列宽
        self.table_view.resizeColumnsToContents()
        
        self.statusBar().showMessage('多级表头已设置：基本信息(个人+基础) | 成绩(文科+理科)')


def main():
    """主函数"""
    app = QApplication(sys.argv)
    window = MultiHeaderDemo()
    window.show()
    sys.exit(app.exec())


if __name__ == '__main__':
    main()
